import toml
import random
import sys

config = toml.load(sys.argv[1])

num_flows = config['num_flows']

num_hosts = config['num_hosts']
num_leaf_switches = config['num_leaf_switches']
num_spine_switches = config['num_spine_switches']
num_switches = num_leaf_switches + num_spine_switches

def get_host_name(i: int) -> str:
    return f'h{i}'

def get_leaf_name(i: int) -> str:
    return f's{i}'

def get_spine_name(i: int) -> str:
    return f's{i + num_leaf_switches}'


ratio = num_hosts// num_leaf_switches
def get_leaf_by_host(i: int) -> int:
    return i // ratio

link_properties = [
    {
        'name': 'default',
        **config['link_property']
    }
]

topo = []

for i in range(num_hosts):
    leaf = get_leaf_name(get_leaf_by_host(i))
    host = get_host_name(i)
    topo.append({
        'link': [leaf, host]
    })

for i in range(num_leaf_switches):
    leaf = get_leaf_name(i)
    for j in range(num_spine_switches):
        spine = get_spine_name(j)
        topo.append({
            'link': [leaf, spine]
        })

flows = []


leaf_switches = list(range(num_leaf_switches))
for i in range(num_flows):
    random.shuffle(leaf_switches)
    leaf1, leaf2 = leaf_switches[:2]
    host1 = ratio * leaf1 + random.randint(0, ratio-1)
    host2 = ratio * leaf2 + random.randint(0, ratio-1)
    spine = random.randint(0, num_spine_switches-1)

    host1 = get_host_name(host1)
    host2 = get_host_name(host2)
    leaf1 = get_leaf_name(leaf1)
    leaf2 = get_leaf_name(leaf2)
    spine = get_spine_name(spine)

    flows.append({
        'path': [host1, leaf1, spine, leaf2, host2]
    })

config = {
    'num_hosts': num_hosts,
    'num_switches': num_switches,
    'link_properties': link_properties,
    'topo': topo,
    'flows': flows,
}

with open(sys.argv[2], 'w') as f:
    toml.dump(config, f)